import os
import os.path as osp
import argparse
from pathlib import Path
from mcp.server.fastmcp import FastMCP

# 创建 MCP Server
mcp = FastMCP("桌面 TXT 文件统计器")
@mcp.tool(
    description="获取指定路径下的所有文件和目录的详细列表。返回结果会通过 [FILE] 和 [DIR] 前缀明确区分文件类型。该工具对于了解目录结构及查找特定文件非常实用。仅限在允许访问的目录中操作。"
)
async def list_directory(dir_path) -> str:
    if not osp.exists(dir_path) or not osp.isdir(dir_path):
        raise ValueError(f"{dir_path} is not a valid directory")

    entries = os.listdir(dir_path)
    formatted = []
    for entry in entries:
        entry_path = osp.join(dir_path, entry)
        if osp.isdir(entry_path):
            formatted.append(f"[DIR] {entry}")
        else:
            formatted.append(f"[FILE] {entry}")

    return "\n".join(formatted)

@mcp.tool()
def count_desktop_txt_files() -> int:
    """Count the number of .txt files on the desktop."""
    # Get the desktop path
    username = os.getenv("USER") or os.getenv("USERNAME")
    desktop_path = Path(f"/Users/{username}/Desktop")

    # Count .txt files
    txt_files = list(desktop_path.glob("*.txt"))
    return len(txt_files)

@mcp.tool()
def list_desktop_txt_files() -> str:
    """Get a list of all .txt filenames on the desktop."""
    # Get the desktop path
    username = os.getenv("USER") or os.getenv("USERNAME")
    desktop_path = Path(f"/Users/{username}/Desktop")

    # Get all .txt files
    txt_files = list(desktop_path.glob("*.txt"))

    # Return the filenames
    if not txt_files:
        return "No .txt files found on desktop."

    # Format the list of filenames
    file_list = "\n".join([f"- {file.name}" for file in txt_files])
    return f"Found {len(txt_files)} .txt files on desktop:\n{file_list}"

if __name__ == "__main__":
    # Initialize and run the server
    mcp.run()